/* main.cpp 
-----------------------------------------
By: 	Wei YE ( ye0003ei@e.ntu.edu.sg )

Date:	03/Oct/2012
-----------------------------------------
*/

#include <highgui.h>
#include <algorithm>
#include "HF.h"
#include "features.h"
#include "MeanShift\ImageHandler.h"
#include "CRF.h"

const int           STEP_WIDTH    = 1;
const int			SEARCH_RADIUS = 20;
const float			WEIGHT_THRESH = 0.3f;
const unsigned int  TREE_NUMBER   = 20;
const unsigned int  TEST_NUMBER   = 8;
const float			backProjectRadius  = 6;
const float			backProjectMinProb = 0.5;


inline Rect getNeighborRegion(const Rect& box, const Size& bound)
{
	Rect nr = Rect(box.x - SEARCH_RADIUS, box.y - SEARCH_RADIUS, box.width  + 2*SEARCH_RADIUS, box.height + 2*SEARCH_RADIUS);
    nr &= Rect(0, 0, bound.width, bound.height);
	return nr;
}

void buildptPosProb(const Mat& img, const Rect& box, const Rect& n_box, float* ptPosProb)
{
	int i, j, k, regionNum;
	ImageHandler handler(8, 4, 100);
	handler.LoadTheImage(img.clone(), COLOR);
	handler.Segment();
	int* labels_tmp = handler.GetLabel();
	if(!labels_tmp) return;

	vector<int> labels(labels_tmp, labels_tmp + img.rows*img.cols);
	regionNum = labels.at(argmax(labels)) + 1;
	
	float* ptCount  = new float[regionNum]; memset(ptCount, 0, regionNum*sizeof(float));
	float* rgWeight = new float[regionNum]; memset(rgWeight,0, regionNum*sizeof(float));
	for(i = 0; i < img.rows; i++){
		for(j = 0; j < img.cols; j++){
			ptCount[labels.at(i*img.cols + j)]++;
			if(i > box.y - n_box.y && i < box.y - n_box.y + box.height 
				&& j > box.x - n_box.x && j < box.x - n_box.x + box.width)
				rgWeight[labels.at(i*img.cols + j)]++;
		}
	}

	for(k = 0; k < regionNum; k++){
		rgWeight[k] /= ptCount[k];
	}

	for(i = 0; i < img.rows; i++){
		for(j = 0; j < img.cols; j++){
			ptPosProb[i*img.cols + j] = rgWeight[labels.at(i*img.cols + j)];
		}
	}

	delete []ptCount;
	delete []rgWeight;
}

void showHeatMap(int height, int width, float* ptPosProb)
{
	Mat hsvMap, heatMap;
	hsvMap = Mat(height, width, CV_8UC3);
	for(int i = 0; i < height; i++){
		for(int j = 0; j < width; j++){
			if(ptPosProb[i*width + j] < 0.5) ptPosProb[i*width + j] = 0;
			hsvMap.at<Vec3b>(i,j) = Vec3b((uchar)(120.0f - ptPosProb[i*width + j]*120.f) , 255, 255);
		}
	}
	cvtColor(hsvMap, heatMap, CV_HSV2BGR);
	namedWindow("HeatMap");
	imshow("HeatMap", heatMap);
	waitKey(1000);
	destroyWindow("HeatMap");
}

void drawResult(Mat& nbImg, int* label, vector<int>& ptLabels, Rect& newBox)
{
	int pr = 0, pb = 0, pl = nbImg.cols, pa = nbImg.rows;
		for(int i = 0; i < nbImg.rows - 1; i++){
			for(int j = 0; j < nbImg.cols - 1; j++){
				if(label[ptLabels.at(i*nbImg.cols + j)] == 0){
					if(label[ptLabels.at(i*nbImg.cols + j + 1)] == 1 || label[ptLabels.at(i*nbImg.cols + j + 1)] == 1)
						nbImg.at<Vec3b>(i,j) = Vec3b(0,0,255);
				}else if(label[ptLabels.at(i*nbImg.cols + j)] == 1)
				{
					if( j > pr) pr = j;
					if( j < pl) pl = j;
					if( i > pb) pb = i;
					if( i < pa) pa = i;
					if(label[ptLabels.at(i*nbImg.cols + j + 1)] == 0 || label[ptLabels.at(i*nbImg.cols + j + 1)] == 0)
						nbImg.at<Vec3b>(i,j) = Vec3b(0,0,255);
					else
						nbImg.at<Vec3b>(i,j)[2] = (uchar)(nbImg.at<Vec3b>(i,j)[2]*0.7f + 255*0.3f);
				}
			}
		}
		newBox = Rect(pl, pa, pr - pl, pb - pa);
}

int main()
{
	// Parameters
	string   dataPath, imgFormat;  
	int      frameStart, frameEnd;
	int      nDigits;  // eg, if nDigits=5, 1->"00001", 12->"00012", likewise.
	Rect     box;

	// Varialbes
	int      i, j, k, lb, cnt;
	int*     label = NULL;
	float*   spPosProb; 
	float*   ptPosProb;
	Rect     nbBox;
	Point    center, minLoc, maxLoc;
	Mat      frame, nbImg, result, out;
	double   minVal, maxVal = 6.0f;

	// Initialization
	ifstream cfg;
	cfg.open("config.txt",ios::in);

	if (!cfg){
		cout<<"Configure file not found!"<<endl;
		return -1;
	}

	cfg>>dataPath; cfg>>imgFormat; cfg>>frameStart; cfg>>frameEnd; cfg>>nDigits;
	if(nDigits != 0)
		assert(pow(10.0f,nDigits) >= frameEnd); // 10^n_digits should >= frameEnd

	cfg>>box.x>>box.y>>box.width>>box.height;
	cfg.close();

	// Read first frame
	frame  = imread(dataPath + int2str(frameStart,nDigits) + imgFormat);
	nbBox  = getNeighborRegion(box, frame.size());
	nbImg  = frame(nbBox);
	center = Point(box.x + box.width/2, box.y + box.height/2); // center in the WHOLE image
	result = Mat(frame.size(), CV_32FC1, Scalar::all(0.0f));

	//-------------------------------------------//
	// USING MEAN SHIFT TO DETERMINE PT_POS_PROB //
	//-------------------------------------------//
	ptPosProb = new float[nbImg.rows*nbImg.cols]; 
	memset(ptPosProb, 0, nbImg.rows*nbImg.cols*sizeof(float));
	buildptPosProb(nbImg, box, nbBox, ptPosProb);
	showHeatMap(nbImg.rows, nbImg.cols, ptPosProb);

	Features ft;
	ft.ExtractFromImage(nbImg);

    //ft.Record(string("C:\\Users\\YE0003EI\\Desktop\\record.txt"));
	spPosProb = new float[ft.SuperpixelNum()]; 
	memset(spPosProb, 0, ft.SuperpixelNum()*sizeof(float));
	vector<int>* pixCount = ft.SuperpixelSize();
	for(i = 0; i < nbImg.rows; i++)
		for(j = 0; j < nbImg.cols; j++)
			spPosProb[ft.PtLabels()->at(i*nbImg.cols + j)] += ptPosProb[i*nbImg.cols + j];

	Ferns fn(TREE_NUMBER, TEST_NUMBER, ft);
	for(i = 0; i < ft.SuperpixelNum(); i++){
		lb = (spPosProb[i]/pixCount->at(i) > 0.5);
		k = poisson(spPosProb[i]/pixCount->at(i));
		for(j = 0; j < k + 5; j++)
			fn.update(ft, i, lb, Point(center.x - nbBox.x, center.y - nbBox.y));
	}	

	putText(frame,("#" + int2str(frameStart, 3)), cv::Point(1, 25), FONT_HERSHEY_SIMPLEX, 1, Scalar(226, 43, 138) , 2);
	rectangle(frame, box, Scalar(226, 43, 138) , 2, 8, 0 );
	namedWindow("tracking");
	imshow("tracking", frame);
	waitKey(30);

	for(int fid = frameStart + 3; fid < frameEnd; fid++)
	{
		frame = imread(dataPath + int2str(fid,nDigits) + imgFormat);
		nbBox = getNeighborRegion(box, frame.size());
		nbImg = Mat(frame, nbBox );
		ft.ExtractFromImage(nbImg);
		fn.evaluate(ft, result, Point(nbBox.x, nbBox.y), STEP_WIDTH, 0.5f);
		out = result;
		float maxv = 0.0f;
		for(int rnm = 0; rnm < result.rows; rnm++){
			for(int cnm = 0; cnm < result.cols; cnm++){
				if(result.at<float>(rnm,cnm) > maxv)
					maxv = result.at<float>(rnm,cnm);
			}
		}
		result *= (1.0f/maxv);
		imshow("tracking", result);
		waitKey(1000);

		normalize(out, out, 255, 0, NORM_MINMAX);
		minMaxLoc(result, &minVal, &maxVal, &minLoc, &maxLoc);
		center = maxLoc;

		if(spPosProb) delete[] spPosProb;
		spPosProb = new float[ft.SuperpixelNum()]; 
		memset(spPosProb, 0, ft.SuperpixelNum()*sizeof(float));
		cnt = fn.backProject(ft, spPosProb, Point(center.x - nbBox.x, center.y - nbBox.y), backProjectRadius, STEP_WIDTH, backProjectMinProb);
		//printArray(spPosProb, ft.SuperpixelNum());
		if(maxVal < 3.0f){
			continue;
		}
		
		if(label) delete[] label;
		label = new int[ft.SuperpixelNum()];
		memset(label, 0, ft.SuperpixelNum()*sizeof(int));
		for(i = 0; i < ft.SuperpixelNum(); i++)
			label[i] = ( spPosProb[i] > 0 );
		//OptLabel(ft, spPosProb, label);
		//printArray(label, ft.SuperpixelNum());

		result = Scalar(0.0f);	
		for(i = 0; i < ft.SuperpixelNum(); i++)
			fn.update(ft, i, label[i], Point(center.x - nbBox.x, center.y - nbBox.y));
		fn.forget(0.90);

		putText(frame,("#" + int2str(frameStart, 3)), cv::Point(1, 25), FONT_HERSHEY_SIMPLEX, 1, Scalar(226, 43, 138) , 2);
		drawResult(nbImg, label, *ft.PtLabels(), box);
		box.x += nbBox.x;
		box.y += nbBox.y;
		imshow("tracking", frame);
		waitKey(30);
	}

	destroyAllWindows();
	delete[] ptPosProb;
	delete[] spPosProb;
	return 0;
}